{
    "cells": [
        {
            "cell_type": "code",
            "execution_count": 1,
            "id": "8a146408",
            "metadata": {},
            "outputs": [],
            "source": [
                "# NBVAL_SKIP\n",
                "import sys\n",
                "sys.path.append('../')\n",
                "\n",
                "import logging\n",
                "logging.getLogger('matplotlib').setLevel(logging.CRITICAL)\n",
                "logging.getLogger('graphein').setLevel(logging.INFO)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "59a32cc0-8bd2-4910-9cb6-980dd2476ff0",
            "metadata": {},
            "source": [
                "# PSCDB - Baselines"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "id": "4ee700f8",
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "RDKit WARNING: [23:36:27] Enabling RDKit 2019.09.3 jupyter extensions\n",
                        "[23:36:27] Enabling RDKit 2019.09.3 jupyter extensions\n"
                    ]
                }
            ],
            "source": [
                "# NBVAL_SKIP\n",
                "import pandas as pd\n",
                "import numpy as np\n",
                "\n",
                "import torch\n",
                "import torch.nn as nn\n",
                "import pytorch_lightning as pl\n",
                "from tqdm.notebook import tqdm\n",
                "import networkx as nx\n",
                "import torch_geometric\n",
                "from torch_geometric.data import Data\n",
                "from torch_geometric.utils import from_networkx\n",
                "from sklearn.preprocessing import LabelBinarizer\n",
                "from sklearn.metrics import f1_score\n",
                "\n",
                "import warnings\n",
                "warnings.filterwarnings(\"ignore\")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "80b5801a",
            "metadata": {},
            "source": [
                "## Load dataset\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "id": "44af34a7",
            "metadata": {},
            "outputs": [],
            "source": [
                "# NBVAL_SKIP\n",
                "df = pd.read_csv(\"../datasets/pscdb/structural_rearrangement_data.csv\")\n",
                "pdbs = df[\"Free PDB\"]\n",
                "y = [torch.argmax(torch.Tensor(lab)).type(torch.LongTensor) for lab in LabelBinarizer().fit_transform(df.motion_type)]"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "595ee7c4",
            "metadata": {},
            "source": [
                "## Transformation from Raw Structure to ML-ready Datasets Construction with Graphein"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "id": "65ac43fd",
            "metadata": {
                "scrolled": true
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "{'granularity': 'CA', 'keep_hets': False, 'insertions': False, 'pdb_dir': PosixPath('../examples/pdbs'), 'verbose': False, 'exclude_waters': True, 'deprotonate': False, 'protein_df_processing_functions': None, 'edge_construction_functions': [functools.partial(<function add_k_nn_edges at 0x7fb4b22bbf70>, k=3, long_interaction_threshold=0)], 'node_metadata_functions': [<function meiler_embedding at 0x7fb4b22c7430>], 'edge_metadata_functions': None, 'graph_metadata_functions': None, 'get_contacts_config': None, 'dssp_config': None}\n"
                    ]
                },
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "ff91e6386f7c422ba50479795794db33",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "  0%|          | 0/891 [00:00<?, ?it/s]"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "URL Error [Errno 110] Connection timed out\n",
                        "274 processing error...\n",
                        "666 processing error...\n",
                        "677 processing error...\n"
                    ]
                }
            ],
            "source": [
                "# NBVAL_SKIP\n",
                "from graphein.protein.config import ProteinGraphConfig\n",
                "from graphein.protein.edges.distance import add_hydrogen_bond_interactions, add_peptide_bonds, add_k_nn_edges\n",
                "from graphein.protein.graphs import construct_graph\n",
                "\n",
                "from functools import partial\n",
                "\n",
                "# Override config with constructors\n",
                "constructors = {\n",
                "    \"edge_construction_functions\": [partial(add_k_nn_edges, k=3, long_interaction_threshold=0)],\n",
                "    #\"edge_construction_functions\": [add_hydrogen_bond_interactions, add_peptide_bonds],\n",
                "    #\"node_metadata_functions\": [add_dssp_feature]\n",
                "}\n",
                "\n",
                "config = ProteinGraphConfig(**constructors)\n",
                "print(config.dict())\n",
                "\n",
                "# Make graphs\n",
                "graph_list = []\n",
                "y_list = []\n",
                "for idx, pdb in enumerate(tqdm(pdbs)):\n",
                "    try:\n",
                "        graph_list.append(\n",
                "            construct_graph(pdb_code=pdb,\n",
                "                        config=config\n",
                "                       )\n",
                "            )\n",
                "        y_list.append(y[idx])\n",
                "    except:\n",
                "        print(str(idx) + ' processing error...')\n",
                "        pass"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 44,
            "id": "eb080ee0-555e-411c-a5e5-8539c97e998f",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "'3e59'"
                        ]
                    },
                    "execution_count": 44,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "# NBVAL_SKIP\n",
                "pdbs[274]\n",
                "#pdbs[266]\n",
                "#pdbs[677]"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "4de24288",
            "metadata": {},
            "source": [
                "## Convert Nx graphs to PyTorch Geometric"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "id": "cd343538",
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Using backend: pytorch\n"
                    ]
                }
            ],
            "source": [
                "# NBVAL_SKIP\n",
                "from graphein.ml.conversion import GraphFormatConvertor\n",
                "\n",
                "format_convertor = GraphFormatConvertor('nx', 'pyg', \n",
                "                                        verbose = 'gnn', \n",
                "                                        columns = None)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "id": "a0f50771",
            "metadata": {
                "tags": []
            },
            "outputs": [
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "563bbc190d4944c9af056dce12eddd39",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "  0%|          | 0/888 [00:00<?, ?it/s]"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                }
            ],
            "source": [
                "# NBVAL_SKIP\n",
                "pyg_list = [format_convertor(graph) for graph in tqdm(graph_list)]"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "id": "79c6eb2e",
            "metadata": {},
            "outputs": [],
            "source": [
                "# NBVAL_SKIP\n",
                "for idx, g in enumerate(pyg_list):\n",
                "    g.y = y_list[idx] \n",
                "    g.coords = torch.FloatTensor(g.coords[0])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "id": "642b78e3",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Data(coords=[10112, 3], dist_mat=[1], edge_index=[2, 120], name=[1], node_id=[1264], y=1)\n",
                        "Data(coords=[820, 3], dist_mat=[1], edge_index=[2, 1431], name=[1], node_id=[808], y=2)\n",
                        "Data(coords=[668, 3], dist_mat=[1], edge_index=[2, 1166], name=[1], node_id=[666], y=4)\n",
                        "Data(coords=[2720, 3], dist_mat=[1], edge_index=[2, 3], name=[1], node_id=[340], y=5)\n"
                    ]
                }
            ],
            "source": [
                "# NBVAL_SKIP\n",
                "for i in pyg_list:\n",
                "    if i.coords.shape[0] == len(i.node_id):\n",
                "        pass\n",
                "    else:\n",
                "        print(i)\n",
                "        pyg_list.remove(i)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "0d07dd30",
            "metadata": {},
            "source": [
                "## Model Configuration"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 32,
            "id": "e1ab2604",
            "metadata": {},
            "outputs": [],
            "source": [
                "# NBVAL_SKIP\n",
                "config_default = dict(\n",
                "    n_hid = 8,\n",
                "    n_out = 8,\n",
                "    batch_size = 4,\n",
                "    dropout = 0.5,\n",
                "    lr = 0.001,\n",
                "    num_heads = 32,\n",
                "    num_att_dim = 64,\n",
                "    model_name = 'GCN'\n",
                ")\n",
                "\n",
                "class Struct:\n",
                "    def __init__(self, **entries):\n",
                "        self.__dict__.update(entries)\n",
                "        \n",
                "config = Struct(**config_default)\n",
                "\n",
                "global model_name\n",
                "model_name = config.model_name"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "14cec1f0",
            "metadata": {},
            "source": [
                "## Construct DataLoaders"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 33,
            "id": "1d67359d",
            "metadata": {},
            "outputs": [],
            "source": [
                "# NBVAL_SKIP\n",
                "import numpy as np\n",
                "np.random.seed(42)\n",
                "idx_all = np.arange(len(pyg_list))\n",
                "np.random.shuffle(idx_all)\n",
                "\n",
                "train_idx, valid_idx, test_idx = np.split(idx_all, [int(.8*len(pyg_list)), int(.9*len(pyg_list))])\n",
                "train, valid, test = [pyg_list[i] for i in train_idx], [pyg_list[i] for i in valid_idx], [pyg_list[i] for i in test_idx]\n",
                "\n",
                "from torch_geometric.data import DataLoader\n",
                "train_loader = DataLoader(train, batch_size=config.batch_size, shuffle = True, drop_last = True)\n",
                "valid_loader = DataLoader(valid, batch_size=32)\n",
                "test_loader = DataLoader(test, batch_size=32)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 34,
            "id": "a8d4193f",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "Data(coords=[635, 3], dist_mat=[1], edge_index=[2, 1118], name=[1], node_id=[635], y=1)"
                        ]
                    },
                    "execution_count": 34,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "# NBVAL_SKIP\n",
                "pyg_list[0]"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "452acd73",
            "metadata": {},
            "source": [
                "## Define Model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 35,
            "id": "53bd0370",
            "metadata": {},
            "outputs": [],
            "source": [
                "# NBVAL_SKIP\n",
                "from torch_geometric.nn import GCNConv, GATConv, SAGEConv, global_add_pool\n",
                "from torch.nn.functional import mse_loss, nll_loss, relu, softmax, cross_entropy\n",
                "from torch.nn import functional as F\n",
                "from pytorch_lightning.metrics.functional import accuracy"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 36,
            "id": "b6728d4f",
            "metadata": {},
            "outputs": [],
            "source": [
                "# NBVAL_SKIP\n",
                "class GraphNets(pl.LightningModule):\n",
                "    def __init__(self):\n",
                "        super().__init__()\n",
                "        \n",
                "        if model_name == 'GCN':\n",
                "            self.layer1 = GCNConv(in_channels=3, out_channels=config.n_hid)\n",
                "            self.layer2 = GCNConv(in_channels=config.n_hid, out_channels=config.n_out)\n",
                "\n",
                "        elif model_name == 'GAT':\n",
                "            self.layer1 = GATConv(3, config.num_att_dim, heads=config.num_heads, dropout=config.dropout)\n",
                "            self.layer2 = GATConv(config.num_att_dim * config.num_heads, out_channels = config.n_out, heads=1, concat=False,\n",
                "                                 dropout=config.dropout)\n",
                "            \n",
                "        elif model_name == 'GraphSAGE':\n",
                "            self.layer1 = SAGEConv(3, config.n_hid)\n",
                "            self.layer2 = SAGEConv(config.n_hid, config.n_out)  \n",
                "            \n",
                "        self.decoder = nn.Linear(config.n_out, 7)\n",
                "        \n",
                "    def forward(self, g):\n",
                "        x = g.coords\n",
                "        x = F.dropout(x, p=config.dropout, training=self.training)\n",
                "        x = F.elu(self.layer1(x, g.edge_index))\n",
                "        x = F.dropout(x, p=config.dropout, training=self.training)\n",
                "        x = self.layer2(x, g.edge_index)\n",
                "        x = global_add_pool(x, batch=g.batch)\n",
                "        x = self.decoder(x)\n",
                "        return softmax(x)\n",
                "\n",
                "    def training_step(self, batch, batch_idx):\n",
                "        x = batch   \n",
                "        y = x.y\n",
                "        y_hat = self(x)\n",
                "        loss = cross_entropy(y_hat, y)\n",
                "        acc = accuracy(y_hat, y)\n",
                "\n",
                "        self.log(\"train_loss\", loss)\n",
                "        self.log(\"train_acc\", acc)\n",
                "        return loss\n",
                "\n",
                "    def validation_step(self, batch, batch_idx):\n",
                "        x = batch   \n",
                "        y = x.y\n",
                "        y_hat = self(x)\n",
                "        loss = cross_entropy(y_hat, y)\n",
                "        acc = accuracy(y_hat, y)\n",
                "        self.log(\"valid_loss\", loss)\n",
                "        self.log(\"valid_acc\", acc)\n",
                "\n",
                "    def test_step(self, batch, batch_idx):\n",
                "        x = batch   \n",
                "        y = x.y\n",
                "        y_hat = self(x)\n",
                "        loss = cross_entropy(y_hat, y)\n",
                "        acc = accuracy(y_hat, y)\n",
                "\n",
                "        y_pred_softmax = torch.log_softmax(y_hat, dim = 1)\n",
                "        y_pred_tags = torch.argmax(y_pred_softmax, dim = 1) \n",
                "        f1 = f1_score(y.detach().cpu().numpy(), y_pred_tags.detach().cpu().numpy(), average = 'weighted')\n",
                "\n",
                "        self.log(\"test_loss\", loss)\n",
                "        self.log(\"test_acc\", acc)\n",
                "        self.log(\"test_f1\", f1)\n",
                "\n",
                "    def configure_optimizers(self):\n",
                "        optimizer = torch.optim.Adam(self.parameters(), lr=config.lr)\n",
                "        return optimizer"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 37,
            "id": "970baced-e3c0-4a17-b0b9-36f5a5a8045d",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "GraphNets(\n",
                            "  (layer1): GCNConv(3, 8)\n",
                            "  (layer2): GCNConv(8, 8)\n",
                            "  (decoder): Linear(in_features=8, out_features=7, bias=True)\n",
                            ")"
                        ]
                    },
                    "execution_count": 37,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "# NBVAL_SKIP\n",
                "GraphNets()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 40,
            "id": "42bc7e82",
            "metadata": {},
            "outputs": [],
            "source": [
                "# NBVAL_SKIP\n",
                "from pytorch_lightning.callbacks import ModelCheckpoint\n",
                "import os\n",
                "\n",
                "file_path = './graphein_model'\n",
                "if not os.path.exists(file_path):\n",
                "    os.mkdir(file_path)\n",
                "\n",
                "checkpoint_callback = ModelCheckpoint(\n",
                "    monitor=\"valid_loss\",\n",
                "    dirpath=file_path,\n",
                "    filename=\"model-{epoch:02d}-{val_loss:.2f}\",\n",
                "    save_top_k=1,\n",
                "    mode=\"min\",\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "3e552f8a",
            "metadata": {},
            "source": [
                "## Train!"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 41,
            "id": "6e14192c",
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "GPU available: True, used: True\n",
                        "TPU available: False, using: 0 TPU cores\n",
                        "IPU available: False, using: 0 IPUs\n",
                        "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
                        "Traceback (most recent call last):\n",
                        "  File \"<string>\", line 1, in <module>\n",
                        "  File \"/home/atj39/anaconda3/envs/graphein-dev/lib/python3.8/multiprocessing/spawn.py\", line 116, in spawn_main\n",
                        "    exitcode = _main(fd, parent_sentinel)\n",
                        "  File \"/home/atj39/anaconda3/envs/graphein-dev/lib/python3.8/multiprocessing/spawn.py\", line 126, in _main\n",
                        "    self = reduction.pickle.load(from_parent)\n",
                        "AttributeError: Can't get attribute 'GraphNets' on <module '__main__' (built-in)>\n"
                    ]
                },
                {
                    "ename": "KeyboardInterrupt",
                    "evalue": "",
                    "output_type": "error",
                    "traceback": [
                        "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
                        "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
                        "\u001b[0;32m/tmp/ipykernel_3426530/3685191286.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mGraphNets\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mtrainer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpl\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTrainer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmax_epochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m200\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgpus\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcallbacks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mcheckpoint_callback\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalid_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;31m# evaluate on the model with the best validation set\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, train_dataloader)\u001b[0m\n\u001b[1;32m    551\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcheckpoint_connector\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mresume_start\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    552\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 553\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    554\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    555\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstopped\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36m_run\u001b[0;34m(self, model)\u001b[0m\n\u001b[1;32m    916\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    917\u001b[0m         \u001b[0;31m# dispatch `start_training` or `start_evaluating` or `start_predicting`\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 918\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dispatch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    919\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    920\u001b[0m         \u001b[0;31m# plugin will finalized fitting (e.g. ddp_spawn will load trained model)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36m_dispatch\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    984\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maccelerator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart_predicting\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    985\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 986\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maccelerator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart_training\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    987\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    988\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mrun_stage\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py\u001b[0m in \u001b[0;36mstart_training\u001b[0;34m(self, trainer)\u001b[0m\n\u001b[1;32m     90\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mstart_training\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m\"pl.Trainer\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 92\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining_type_plugin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart_training\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     93\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     94\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mstart_evaluating\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m\"pl.Trainer\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp_spawn.py\u001b[0m in \u001b[0;36mstart_training\u001b[0;34m(self, trainer)\u001b[0m\n\u001b[1;32m    156\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    157\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mstart_training\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 158\u001b[0;31m         \u001b[0mmp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mspawn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnew_process\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmp_spawn_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    159\u001b[0m         \u001b[0;31m# reset optimizers, since main process is never used for training and thus does not have a valid optim state\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    160\u001b[0m         \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/site-packages/torch/multiprocessing/spawn.py\u001b[0m in \u001b[0;36mspawn\u001b[0;34m(fn, args, nprocs, join, daemon, start_method)\u001b[0m\n\u001b[1;32m    228\u001b[0m                ' torch.multiprocessing.start_processes(...)' % start_method)\n\u001b[1;32m    229\u001b[0m         \u001b[0mwarnings\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwarn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 230\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mstart_processes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnprocs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mjoin\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdaemon\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_method\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'spawn'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/site-packages/torch/multiprocessing/spawn.py\u001b[0m in \u001b[0;36mstart_processes\u001b[0;34m(fn, args, nprocs, join, daemon, start_method)\u001b[0m\n\u001b[1;32m    177\u001b[0m             \u001b[0mdaemon\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdaemon\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    178\u001b[0m         )\n\u001b[0;32m--> 179\u001b[0;31m         \u001b[0mprocess\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    180\u001b[0m         \u001b[0merror_queues\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0merror_queue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    181\u001b[0m         \u001b[0mprocesses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprocess\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/multiprocessing/process.py\u001b[0m in \u001b[0;36mstart\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    119\u001b[0m                \u001b[0;34m'daemonic processes are not allowed to have children'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    120\u001b[0m         \u001b[0m_cleanup\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 121\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_popen\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_Popen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    122\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_sentinel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_popen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msentinel\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    123\u001b[0m         \u001b[0;31m# Avoid a refcycle if the target function holds an indirect\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/multiprocessing/context.py\u001b[0m in \u001b[0;36m_Popen\u001b[0;34m(process_obj)\u001b[0m\n\u001b[1;32m    282\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0m_Popen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprocess_obj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    283\u001b[0m             \u001b[0;32mfrom\u001b[0m \u001b[0;34m.\u001b[0m\u001b[0mpopen_spawn_posix\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mPopen\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 284\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mPopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprocess_obj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    285\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    286\u001b[0m     \u001b[0;32mclass\u001b[0m \u001b[0mForkServerProcess\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprocess\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mBaseProcess\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/multiprocessing/popen_spawn_posix.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, process_obj)\u001b[0m\n\u001b[1;32m     30\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprocess_obj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     31\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_fds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 32\u001b[0;31m         \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprocess_obj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     34\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mduplicate_for_child\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfd\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/multiprocessing/popen_fork.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, process_obj)\u001b[0m\n\u001b[1;32m     17\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreturncode\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfinalizer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 19\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_launch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprocess_obj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     20\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mduplicate_for_child\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfd\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/anaconda3/envs/graphein-dev/lib/python3.8/multiprocessing/popen_spawn_posix.py\u001b[0m in \u001b[0;36m_launch\u001b[0;34m(self, process_obj)\u001b[0m\n\u001b[1;32m     60\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msentinel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mparent_r\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     61\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparent_w\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'wb'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mclosefd\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 62\u001b[0;31m                 \u001b[0mf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwrite\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgetbuffer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     63\u001b[0m         \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     64\u001b[0m             \u001b[0mfds_to_close\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
                    ]
                }
            ],
            "source": [
                "# NBVAL_SKIP\n",
                "# Train Model\n",
                "model = GraphNets()\n",
                "trainer = pl.Trainer(max_epochs=200, gpus=-1, callbacks=[checkpoint_callback])\n",
                "trainer.fit(model, train_loader, valid_loader)\n",
                "\n",
                "# evaluate on the model with the best validation set\n",
                "best_model = GraphNets.load_from_checkpoint(checkpoint_callback.best_model_path)\n",
                "out_best_test = trainer.test(best_model, test_loader)[0]"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.8.10"
        },
        "nbsphinx": {
            "execute": "never"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}
